# %%
%matplotlib inline
import matplotlib.pyplot as plt, numpy as np, pims, pandas as pd, cv2 as cv, scipy.optimize as op, seaborn as sns


def ls_func(img):
    # filter image
    img = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
    img = cv.medianBlur(img, 1)
    # apply adaptive threshold
    img = cv.adaptiveThreshold(img, 255, cv.ADAPTIVE_THRESH_GAUSSIAN_C, cv.THRESH_BINARY, 1501, 1)
    radius = 50
    X, Y = np.meshgrid(np.linspace(-radius/2, radius/2, radius+1), np.linspace(-radius/2, radius/2, radius+1))
    kernel = np.where(np.sqrt(X**2 + Y**2)<radius/2, 1, 0).astype(np.uint8)
    img = cv.morphologyEx(cv.bitwise_not(img), cv.MORPH_CLOSE, kernel, iterations=1)
    # reverse closing to get rid of noise above bed
    img = cv.morphologyEx(cv.bitwise_not(img), cv.MORPH_CLOSE, kernel, iterations=1)
    # Extract surface of bead mound
    return cv.findContours(img, cv.RETR_LIST, 2)
    # return img

def oc_func(img):
    # filter image
    img = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
    img = cv.medianBlur(img, 21)
    # apply adaptive threshold
    img = cv.adaptiveThreshold(img, 255, cv.ADAPTIVE_THRESH_GAUSSIAN_C, cv.THRESH_BINARY, 1501, 1)
    radius = 40
    X, Y = np.meshgrid(np.linspace(-radius/2, radius/2, radius+1), np.linspace(-radius/2, radius/2, radius+1))
    kernel = np.where(np.sqrt(X**2 + Y**2)<radius/2, 1, 0).astype(np.uint8)
    # img = cv.morphologyEx(cv.bitwise_not(img), cv.MORPH_CLOSE, kernel, iterations=1)
    # reverse closing to get rid of noise above bed
    # img = cv.morphologyEx(cv.bitwise_not(img), cv.MORPH_CLOSE, kernel, iterations=1)
    # Extract surface of bead mound
    # img = cv.Canny(img,100,100)
    # return img
    return cv.findContours(img, cv.RETR_LIST, 2)

def gg_func(img):
    # filter image
    img = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
    img = cv.medianBlur(img, 1)
    # apply adaptive threshold
    img = cv.adaptiveThreshold(img, 255, cv.ADAPTIVE_THRESH_GAUSSIAN_C, cv.THRESH_BINARY, 2001, 1)
    radius = 50
    X, Y = np.meshgrid(np.linspace(-radius/2, radius/2, radius+1), np.linspace(-radius/2, radius/2, radius+1))
    kernel = np.where(np.sqrt(X**2 + Y**2)<radius/2, 1, 0).astype(np.uint8)
    img = cv.morphologyEx(cv.bitwise_not(img), cv.MORPH_CLOSE, kernel, iterations=1)
    # reverse closing to get rid of noise above bed
    img = cv.morphologyEx(cv.bitwise_not(img), cv.MORPH_CLOSE, kernel, iterations=1)
    # Extract surface of bead mound
    # img = cv.Canny(img,100,100)
    # return img    
    return cv.findContours(img, cv.RETR_LIST, 2)

def ng_func(img):
    # filter image
    # img = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
    # img = cv.medianBlur(img, 1)
    img1 = np.asarray(img[:,:,0])
    img1 = cv.medianBlur(img1, 1)
    # img1 = cv.adaptiveThreshold(img1, 255, cv.ADAPTIVE_THRESH_GAUSSIAN_C, cv.THRESH_BINARY, 1501, 11)
    img2 = np.asarray(img[:,:,1])
    img2 = cv.medianBlur(img2, 1)
    # img2 = cv.adaptiveThreshold(img2, 255, cv.ADAPTIVE_THRESH_GAUSSIAN_C, cv.THRESH_BINARY, 1501, 11)
    img3 = np.asarray(img[:,:,2])
    img3 = cv.medianBlur(img3, 1)
    # img3 = cv.adaptiveThreshold(img3, 255, cv.ADAPTIVE_THRESH_GAUSSIAN_C, cv.THRESH_BINARY, 1501, 11)
    img = np.maximum(img1, img2, img3)
    # apply adaptive threshold
    img = cv.adaptiveThreshold(img, 255, cv.ADAPTIVE_THRESH_GAUSSIAN_C, cv.THRESH_BINARY, 1501, 1)
    radius = 50
    X, Y = np.meshgrid(np.linspace(-radius/2, radius/2, radius+1), np.linspace(-radius/2, radius/2, radius+1))
    kernel = np.where(np.sqrt(X**2 + Y**2)<radius/2, 1, 0).astype(np.uint8)
    img = cv.morphologyEx(cv.bitwise_not(img), cv.MORPH_CLOSE, kernel, iterations=1)
    # reverse closing to get rid of noise above bed
    img = cv.morphologyEx(cv.bitwise_not(img), cv.MORPH_CLOSE, kernel, iterations=1)
    # Extract surface of bead mound
    # img = cv.Canny(img,100,100)
    # return img    
    return cv.findContours(img, cv.RETR_LIST, 2)

grains = ['gg', 'oc', 'ng', 'ngup', 'gb', 'ls']
funcs = {'gg': gg_func, 'ng': ng_func, 'ls': ls_func, 'oc': oc_func}
fns = {'ng': 'shape_images/ng.jpg', 'ngup': 'shape_images/ng_up.jpg', 'gg': 'shape_images/gg.jpg', 'oc': 'shape_images/oc.jpg', 'gb': 'shape_images/gb.jpg', 'ls': 'shape_images/ls.jpg'}

#%%

# iterate over grain types
proc_img = {}
for key in ['ls', 'gg', 'ng', 'oc']:
    print(key)
    if key is 'oc':
        cont, heir = funcs[key](pims.ImageSequence(fns[key])[0][1700:2050,1350:1750,:])
    else:
        cont, heir = funcs[key](pims.ImageSequence(fns[key])[0][:,500:,:])
    proc_img[key] = {'contours': cont, 'heirarchy': heir}

# %%

def stats(cnt):
    area = cv.contourArea(cnt)
    # print(circularity)
    if area > 10000:
        perimeter = cv.arcLength(cnt,True)
        # print(perimeter)
        circularity = 4 * np.pi * area / perimeter**2
        # print(area)
        return [area, perimeter, circularity]
    else:
        return [np.nan, np.nan, np.nan]

for grain in ['ls', 'gg', 'ng', 'oc']:
    stat_list = np.array([stats(cnt) for cnt in proc_img[grain]['contours']])
    sns.distplot(stat_list[:,2])
    print(grain, np.nanmean(stat_list[:,2]), np.nanstd(stat_list[:,2]))


# %%
